
from pydantic import Field
from core.llm.base_analyze import BaseAnalyzer
from core import prompt_dir
from typing import Dict, Any
import pandas as pd
from loguru import logger


class WYXRewriteQueryAnalyzer(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / '问题改写-王宇翔.md')

    def invoke(
            self,
            userinfo: Any = Field(description="用户的信息"),
            query: str = Field(description='用户的问题描述'),
            schemas: Dict = Field(description="数据表结构描述")
    ):

        logger.info(f'userinfo: {userinfo}')
        logger.info(f'query: {query}')
        logger.info(f'schemas: {schemas}')

        short_schemas = []
        for schema in schemas:
            fields = schema.get('fields', [])
            short_schemas.append({
                '表名': schema.get('name', ''),
                '列名描述': [item.get('name') if isinstance(item, dict) else item for item in fields],
            })
        str_schemas = [str(item) for item in short_schemas]

        prompt = self.template_prompt.format(
            inputs={
                "userinfo": userinfo,
                "query": query,
                "schemas": str_schemas
            },
            remove_template_variables=True
        )
        logger.info(f'改写的提示词:\n{prompt}')

        system_prompt = """
您是一个数据分析专家，擅长于将用户提问改写成跟针对某个指标表的提问。
您依据您数据分析的经验，用户提问改写过程是一个逐步推理分析的过程
，具体包含"找出跟用户提问相关的指标表->每个相关指标表生成一个新的问题->汇总出最终提问列表"三个主要过程，举例如下:
'''
用户提问为"我今年科研情况"。

# 用户信息
我是一个教职工，工号为200232

# 指标列表为:
发表论文数(数量(long),工号(string),年份(STRING))
科研项目数(数量(FLOAT),工号(string),年份(STRING))
科研奖金数(数量(FLOAT),年份(STRING))；
教室排课表(排课总数(LONG),教室名称(STRING),节次名称(STRING))

# 改写过程如下:

1. 找出跟用户提问相关的指标表
   发表论文数、科研项目书跟用户提问"我今年科研情况"相关，科研奖金数虽然属于科研情况但因为没有工号列无法查询出我的科研奖金数，所以跟用户提问不相干。

2. 每个相关指标表生成一个新的问题
   发表论文数: 我今年发表论文情况
   科研项目数: 我今年科研项目情况

3. 汇总出最终提问列表
   改写后的提问为:["我今年发表论文情况","我今年科研项目情况"]
'''
        """

        data = self.analyze(prompt, system_prompt)
        logger.info(f'问题改写结果： {data}')

        if isinstance(data, list) and len(data) > 0:
            rewrites = data
        else:
            rewrites = [query]
        return rewrites
